﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;

namespace ArchiveLink
{
    using System.IO;
    using System.Management;

    public class DriveWatcher:IDisposable
    {
        private ManagementEventWatcher DiskEventWatcher { get; set; }
        private readonly object _updateLock = new object();
        private readonly Dictionary<string, DriveEventArgs> _knownDrives = new Dictionary<string, DriveEventArgs>();
        
        public delegate void DriveArrivedHandler(object sender, DriveEventArgs e);
        public event DriveArrivedHandler DriveArrived;
        public event DriveArrivedHandler DriveRemoved;

        [DllImport("kernel32.dll")]
        private static extern long GetVolumeInformation(string PathName, StringBuilder VolumeNameBuffer, UInt32 VolumeNameSize, ref UInt32 VolumeSerialNumber, ref UInt32 MaximumComponentLength, ref UInt32 FileSystemFlags, StringBuilder FileSystemNameBuffer, UInt32 FileSystemNameSize);


        public DriveWatcher()
        {
            UpdateDrives();
        }
        
        private void UpdateDrives()
        {
            lock (_updateLock)
            {
                InternalUpdateDrives();
            }
        }

        private void InternalUpdateDrives()
        {
            var drives = DriveInfo.GetDrives();
            foreach (var drive in drives)
            {
                if (drive.IsReady)
                {
                    var args = new DriveEventArgs(drive.VolumeLabel, drive.Name,
                                                               GetVolumeSerial(drive));
                    AddToKnown(args);
                }
            }
        }
        public static uint GetVolumeSerial(DriveInfo drive)
        {
            uint serNum = 0;
            uint maxCompLen = 0;
            StringBuilder VolLabel = new StringBuilder(256);
            UInt32 VolFlags = new UInt32();
            StringBuilder fileSystemName = new StringBuilder(256);
            long retVal = GetVolumeInformation(drive.Name, VolLabel, (UInt32)VolLabel.Capacity, ref serNum, ref maxCompLen, ref VolFlags, fileSystemName, (UInt32)fileSystemName.Capacity);
            return serNum;
        }

        public void InitDiskEvent(bool considerOpticalDisks)
        {
            WqlEventQuery q;

            // Bind to local machine
            ConnectionOptions opt = new ConnectionOptions();
            opt.EnablePrivileges = true; //sets required privilege
            ManagementScope scope = new ManagementScope("root\\CIMV2", opt);

            try
            {
                q = new WqlEventQuery();
                q.EventClassName = "__InstanceOperationEvent";
                q.WithinInterval = new TimeSpan(0, 0, 1);

                // DriveType - 2: Removable
                // DriveType - 3: Harddisk
                // DriveType - 5: CDROM
                if (considerOpticalDisks)
                {
                    q.Condition = @"TargetInstance ISA 'Win32_LogicalDisk' and (TargetInstance.DriveType = 2 or TargetInstance.DriveType = 3 or TargetInstance.DriveType = 5)";
                }
                else
                {
                    q.Condition = @"TargetInstance ISA 'Win32_LogicalDisk' and (TargetInstance.DriveType = 2 or TargetInstance.DriveType = 3)";                    
                    //q.Condition = @"TargetInstance ISA 'Win32_LogicalDisk' and (TargetInstance.DriveType = 2)";                    
                }
               
                DiskEventWatcher = new ManagementEventWatcher(scope, q);

                // register async. event handler
                DiskEventWatcher.EventArrived += (DiskEventArrived);
                DiskEventWatcher.Start();
            }
            catch (Exception e)
            {
                DiskEventWatcher.Stop();
                throw;
            }
        }

        readonly List<string> waitForMod = new List<string>();

        void DiskEventArrived(object sender, EventArrivedEventArgs e)
        {
            // Get the Event object and display it
            PropertyData pd = e.NewEvent.Properties["TargetInstance"];
            
            if (pd != null)
            {
                ManagementBaseObject mbo = pd.Value as ManagementBaseObject;

                //foreach (var property in mbo.Properties)
                //{
                //    Console.WriteLine(property.Name + ":" + property.Value);
                //}

                //Access:0
                //Availability:
                //BlockSize:
                //Caption:G:
                //Compressed:False
                //ConfigManagerErrorCode:
                //ConfigManagerUserConfig:
                //CreationClassName:Win32_LogicalDisk
                //Description:Removable Disk
                //DeviceID:G:
                //DriveType:2
                //ErrorCleared:
                //ErrorDescription:
                //ErrorMethodology:
                //FileSystem:FAT32
                //FreeSpace:11656003584
                //InstallDate:
                //LastErrorCode:
                //MaximumComponentLength:255
                //MediaType:
                //Name:G:
                //NumberOfBlocks:
                //PNPDeviceID:
                //PowerManagementCapabilities:
                //PowerManagementSupported:
                //ProviderName:
                //Purpose:
                //QuotasDisabled:
                //QuotasIncomplete:
                //QuotasRebuilding:
                //Size:15347220480
                //Status:
                //StatusInfo:
                //SupportsDiskQuotas:False
                //SupportsFileBasedCompression:False
                //SystemCreationClassName:Win32_ComputerSystem
                //SystemName:S4M01908
                //VolumeDirty:False
                //VolumeName:bla
                //VolumeSerialNumber:12028A64

                var driveType = (uint)mbo.Properties["DriveType"].Value;
                var volumeName = (string)mbo.Properties["VolumeName"].Value;
                var volumeSerialHex = (string)mbo.Properties["VolumeSerialNumber"].Value;
                uint volumeSerial=0;
                if (!string.IsNullOrWhiteSpace(volumeSerialHex))
                {
                    volumeSerial = uint.Parse(volumeSerialHex, System.Globalization.NumberStyles.HexNumber);
                }
                var drive = (string)mbo.Properties["Name"].Value;

                switch (driveType)
                {
                    case 2:
                    case 3:
                        switch (e.NewEvent.ClassPath.ClassName)
                        {
                            case "__InstanceCreationEvent":
                                Console.WriteLine("Drive has been inserted");
                                if (volumeSerial==0)
                                {
                                    lock (waitForMod)
                                    {
                                        waitForMod.Add(drive);    
                                    }
                                    break;
                                }
                                var args = new DriveEventArgs(volumeName, drive, volumeSerial);
                                AddToKnown(args);
                                if (DriveArrived != null)
                                {
                                    DriveArrived(this, args);
                                }
                                break;
                            case "__InstanceDeletionEvent":
                                Console.WriteLine("Drive has been ejected");

                                var args3=new DriveEventArgs(volumeName, drive, volumeSerial);
                                args3=RemoveFromKnown(args3);
                                if (args3!=null&&DriveRemoved != null)
                                {
                                    DriveRemoved(this, args3);
                                }
                                break;                         
                            case "__InstanceModificationEvent":
                                lock (waitForMod)
                                {
                                    if (waitForMod.Count == 0)
                                    {
                                        break;
                                    }


                                    if (waitForMod.Contains(drive) && volumeSerial>0)
                                    {
                                        Console.WriteLine("Drive has been inserted");
                                        waitForMod.Remove(drive);

                                        var args2 = new DriveEventArgs(volumeName, drive, volumeSerial);
                                        AddToKnown(args2);
                                        if (DriveArrived != null)
                                        {
                                            DriveArrived(this, args2);
                                        }
                                    }
                                }
                                break;
                        }
                        break;
                    case 5:
                        if (e.NewEvent.ClassPath.ClassName == "__InstanceModificationEvent")
                        {
                            // if CD removed VolumeName == null
                            if (volumeName != null)
                            {
                                Console.WriteLine("Disk has been inserted");

                                var args4 = new DriveEventArgs(volumeName, drive, volumeSerial);
                                AddToKnown(args4);
                                if (DriveArrived != null)
                                {
                                    DriveArrived(this, args4);
                                }
                            }
                            else
                            {
                                Console.WriteLine("Disk has been ejected");
                            }
                        }

                        break;
                }
                
                //UpdateDrives();
            }

        }

        private DriveEventArgs RemoveFromKnown(DriveEventArgs args)
        {
            lock (_knownDrives)
            {
                var driveLetter = args.Drive;
                if (_knownDrives.ContainsKey(driveLetter))
                {
                    args = _knownDrives[driveLetter];
                    _knownDrives.Remove(driveLetter);
                    return args;
                }
                else
                {
                    return null;
                }
            }
        }

        private void AddToKnown(DriveEventArgs args)
        {
            lock (_knownDrives)
            {
                var driveLetter = args.Drive;
                if (_knownDrives.ContainsKey(driveLetter))
                {
                    _knownDrives[driveLetter] = args;
                }
                else
                {
                    _knownDrives.Add(driveLetter, args);
                }
            }
        }

        #region Implementation of IDisposable

        public void Dispose()
        {
            if (DiskEventWatcher != null)
            {
                DiskEventWatcher.Stop();
            }
        }

        #endregion

        public List<DriveEventArgs> GetKnownAndMirrored(List<uint> archiveVolumes)
        {
            lock (_knownDrives)
            {
                return _knownDrives.Values.Where(kd => archiveVolumes.Contains(kd.Serial)).ToList();
            }
        }

        public bool IsConnected(string label)
        {
            lock (_knownDrives)
            {
                return _knownDrives.Values.Any(kd => kd.Name==label);
            }
        }

        public DriveEventArgs GetConnectedInfo(string label)
        {
            lock (_knownDrives)
            {
                return _knownDrives.Values.First(kd => kd.Name == label);
            }
        }
    }
}
